import os
import time
import numpy as np
from sklearn.preprocessing import normalize
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import torch
import torchvision as tv
import torchvision.transforms as transforms
from args import *
from tool import *
import evaluate
# from models import *


save_path = os.path.join(args.log_path, "image.{}".format(args.model_type))
if not os.path.exists(save_path):
    os.makedirs(save_path)


test_loader = torch.utils.data.DataLoader(
    tv.datasets.MNIST(args.data_path, train=False,
                   transform=transforms.Compose([
                    #    transforms.Resize([224, 224]),
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,)),
                    #    transforms.Lambda(lambda x: x.repeat(3, 1, 1))
                   ])),
    batch_size=100, shuffle=True)


P = os.path.join(args.log_path, "am-softmax_0.4m_5.0s")

model = torch.load(os.path.join(P, "19e_0.9916acc.pth")).cuda()

C = np.load(os.path.join(P, "C.19e.npy"))
N_CLASS = C.shape[0]


def get_fea():
    F, L = [], []
    _cnt = 0
    with torch.no_grad():
        for X, Y in test_loader:
            # X, Y = X.cuda(), Y.cuda()
            X = X.cuda()
            _f, _ = model(X)
            F.append(_f.cpu().numpy())
            L.append(Y.numpy())
            _cnt += _f.size(0)
            if _cnt >= 2000:
                break

    F = np.vstack(F)
    L = np.concatenate(L)
    return F, L


F, L = get_fea()
W = model.state_dict()["clf_layer.weight"].cpu().numpy()  # [10, 2]


title = "feature weight vector centre"
fig, ax = plt.subplots(1, 2)
plt.title(title)
ax[0].set_title("features & weight vectors")
ax[1].set_title("features & centres")
F_n = normalize(F, "l2", 1)
W_n = normalize(W, "l2", 1)
C_n = normalize(C, "l2", 1)
ax[0].scatter(F_n[:, 0], F_n[:, 1], s=10, c=L, marker='.', cmap="rainbow")
ax[1].scatter(F_n[:, 0], F_n[:, 1], s=10, c=L, marker='.', cmap="rainbow")
for i in range(W.shape[0]):
    ax[0].plot([0, W_n[i][0]], [0, W_n[i][1]])
    ax[1].plot([0, C_n[i][0]], [0, C_n[i][1]])
fig.savefig(os.path.join(save_path, "{}.png".format(title.replace(" ", "_"))))
plt.close(fig)
